import json
import time
import torch
from celery import shared_task
from core.redis_client import redis_client

@shared_task
def run_task(job_id: str, task_type: str, data_file: str, params: dict = None, model_id: str = None):
    try:
        # 更新状态为running
        redis_client.hset(f"job:{job_id}", mapping={
            "status": "running",
            "startedAt": str(int(time.time()))
        })
        redis_client.publish(f"job-updates:{job_id}", json.dumps({"status": "running"}))

        # 模拟任务执行（替换为实际逻辑）
        if task_type == "Training":
            # 加载数据，训练模型
            for epoch in range(params.get("numEpochs", 5)):
                for step in range(100):  # 假设100步
                    # 检查停止信号
                    if redis_client.hget(f"job:{job_id}", "stopRequested") == "1":
                        redis_client.hset(f"job:{job_id}", "status", "stopped")
                        redis_client.publish(f"job-updates:{job_id}", json.dumps({"status": "stopped"}))
                        return "Stopped"
                    # 模拟进度
                    progress = {"epoch": epoch + 1, "step": step + 1, "totalSteps": 100, "loss": 0.123, "accuracy": 0.95}
                    redis_client.hset(f"job:{job_id}", "progress", json.dumps(progress))
                    redis_client.publish(f"job-updates:{job_id}", json.dumps(progress))
                    time.sleep(0.1)  # 模拟延时

            # 保存模型（示例）
            model_path = f"{settings.MODELS_DIR}/resnext50_{time.strftime('%Y%m%d_%H%M')}.pth"
            # torch.save(model.state_dict(), model_path)

        elif task_type in ["Validation", "Test"]:
            # 加载模型，进行验证/测试
            # 假设结果
            results = [{"name": "类别A", "precision": 0.98, "recall": 0.95, "f1": 0.96}]
            redis_client.hset(f"job:{job_id}", "result", json.dumps(results))

        # 更新完成状态
        redis_client.hset(f"job:{job_id}", mapping={
            "status": "completed",
            "finishedAt": str(int(time.time()))
        })
        redis_client.publish(f"job-updates:{job_id}", json.dumps({"status": "completed"}))

    except Exception as e:
        redis_client.hset(f"job:{job_id}", mapping={
            "status": "failed",
            "error": str(e),
            "finishedAt": str(int(time.time()))
        })
        redis_client.publish(f"job-updates:{job_id}", json.dumps({"status": "failed", "error": str(e)}))